'''
# 柳博存储的suvmax2.csv文件CT和PET目录存在问题，重新整理存储为data.csv
# 转存要用到的CT和PET dicom格式切片，加速随后的图像预处理
# labels.csv文件可视为已知文件，由于最早提取的标签文件格式有问题，所以这里重新生成
# 提取CT和PET中包含的额外信息，生成all_data.csv
# 原始切片中的病人身高单位不统一（cm,m）,故统一单位为 cm
'''


import pandas as pd
import os
import pydicom as pydic
import shutil
import csv
import numpy as np


# 重新整理存储suvmax2.csv
def rebuild_path():
    data = pd.read_csv('D:/lung_cancer/data/suvmax2.csv')

    columns = ['idx', 'patientID', 'z', 'x', 'y', 'r', 'cancer_type', 'CT_SeriesDescription',
               'PET_SeriesDescription', 'origin_dir', 'CT_origin_path', 'PET_origin_path',
               'CT_size', 'PET_size', 'CT_pixel_spacing', 'PET_pixel_spacing', 'patientWeight',
               'patientSex', 'patientAge', 'TotalDose', 't0', 't1', 'HalfLife', 'DecayFactor',
               'CT_slice_path', 'PET_slice_path', 'part_suvmax', 'global_suvmax', 'suv_avg',
               'suv_min', 'suv_std']
    data_list = []

    for i in range(len(data['idx'])):
        length = len(data['origin_dir'][i])
        new_CT_path = data['origin_dir'][i] + '/' + data['CT_OriginPath'][i][length:]
        new_PET_path = data['origin_dir'][i] + '/' + data['PET_OriginPath'][i][length:]
        patientID = data['patient_id'][i]
        z_slice = data['z_slice'][i]
        x_pix = data['x_pix'][i]
        y_pix = data['y_pix'][i]
        ct_r_pix = data['ct_r_pix'][i]
        cancer_type = data['cancer_type'][i]
        CT_SeriesDescription = data['CT_SeriesDescription'][i]
        PET_SeriesDescription = data['PET_SeriesDescription'][i]
        CT_size = data['CT_size'][i]
        PET_size = data['PET_size'][i]
        CT_pixel_spacing = data['CT_pixel_spacing'][i]
        PET_pixel_spacing = data['PET_pixel_spacing'][i]

        slice = pydic.read_file(os.path.join('H:', new_CT_path))
        PatientWeight = slice[0x10, 0x1030].value
        PatientAge = int(slice[0x10, 0x1010].value[:-1])
        PatientSex = slice[0x10, 0x40].value
        Sex = 0  # 男性为0， 女性为1
        if(PatientSex=='F'):
            Sex=1
        TotalDose = data['TotalDose'][i]
        t1 = data['t1'][i]
        t0 = data['t0'][i]
        HalfLife = data['HalfLife'][i]
        DecayFactor = data['DecayFactor'][i]
        CT_slice_path = data['CT_slice_path'][i][19:]
        PET_slice_path = data['PET_slice_path'][i][19:]
        part_suvmax = data['SUVMAX'][i]
        global_suvmax = data['SUVMAX2'][i]
        suv_avg = data['SUVAVG'][i]
        suv_min = data['SUVMIN'][i]
        suv_std = data['SUVSTD'][i]

        data_list.append([i, patientID, z_slice, x_pix, y_pix, ct_r_pix, cancer_type, CT_SeriesDescription,
               PET_SeriesDescription, data['origin_dir'][i], new_CT_path, new_PET_path,
               CT_size, PET_size, CT_pixel_spacing, PET_pixel_spacing, PatientWeight,
               Sex, PatientAge, TotalDose, t0, t1, HalfLife, DecayFactor,
               CT_slice_path, PET_slice_path, part_suvmax, global_suvmax, suv_avg,
               suv_min, suv_std])

    df = pd.DataFrame(data_list, columns=columns)
    df.to_csv('D:/lung_cancer/data/data.csv', index=False)


# 转存使用的数据中的第‘z’张CT和PET
def extract_origin_CT_PET():
    data = pd.read_csv('D:/lung_cancer/data/data.csv')
    CT_path = data['CT_origin_path']
    PET_path = data['PET_origin_path']
    patientID = data['patientID']

    for i in range(len(patientID)):
        CT_save_path = 'D:/lung_cancer/data/origin_data'+'/'+patientID[i]+'/CT/'
        PET_save_path = 'D:/lung_cancer/data/origin_data' + '/' + patientID[i] + '/PET/'
        if not os.path.exists(CT_save_path):
            os.makedirs(CT_save_path)
        if not os.path.exists(PET_save_path):
            os.makedirs(PET_save_path)

        shutil.copyfile('H:/'+CT_path[i], CT_save_path+'/'+CT_path[i].split('/')[-1])
        shutil.copyfile('H:/' + PET_path[i], PET_save_path + '/' + PET_path[i].split('/')[-1])


# 生成labels.csv
def resaved_label():
    columns = ['patientID', 'z', 'x', 'y', 'r', 'cancer_type', 'CT_SeriesDescription',
               'PET_SeriesDescription', 'origin_dir', 'CT_origin_path', 'PET_origin_path']
    data = []
    myfile = open('D:/lung_cancer/data/data.csv')
    lines = csv.reader(myfile)
    for line in lines:
        data.append(line)

    labels = []
    for line in data[1:]:
        labels.append(line[1:12])
    df = pd.DataFrame(labels, columns=columns)
    df.to_csv('D:/lung_cancer/data/labels/labels.csv', index=True)


# 保存CT和PET信息: all_data.csv
def save_info():

    num = 1

    columns = ['patientID', 'z', 'x', 'y', 'r', 'cancer_type', 'ct_size', 'pet_size', 'ct_x_spacing', 'ct_y_spacing',
               'ct_z_spacing', 'pet_x_spacing', 'pet_y_spacing', 'pet_z_spacing', 'ct_intercept', 'pet_intercept',
               'ct_slope', 'pet_slope', 'patientWeight', 'patientSex', 'patientAge', 'patientSize', 'TotalDose', 't0', 't1',
               'HalfLife', 'DecayFactor', 'CT_SeriesDescription', 'PET_SeriesDescription', 'origin_dir',
               'CT_origin_path', 'PET_origin_path']

    myfile = open('D:/lung_cancer/data/labels/labels.csv')
    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []
    for line in labels[1:]:

        path = 'H:/'+line[9]
        names = os.listdir(path)
        ct_slices = []
        pet_slices = []

        for name in names:
            slice = pydic.read_file(path + '/' + name)
            if slice.SeriesDescription == line[7]:
                ct_slices.append(slice)
            elif slice.SeriesDescription == line[8]:
                pet_slices.append(slice)
        ct_slices.sort(key=lambda x: int(x.InstanceNumber))
        pet_slices.sort(key=lambda x: int(x.InstanceNumber))
        ct_z_spacing = np.abs(ct_slices[0].SliceLocation - ct_slices[1].SliceLocation)
        pet_z_spacing = np.abs(pet_slices[0].SliceLocation - pet_slices[1].SliceLocation)


        ct = pydic.read_file('H:/'+line[-2])
        pet = pydic.read_file('H:/'+line[-1])

        patientID = line[1]
        z = line[2]
        x = line[3]
        y = line[4]
        r = line[5]
        cancer_type = line[6]
        ct_size = ct.Rows
        pet_size = pet.Rows
        ct_x_spacing = ct.PixelSpacing[0]
        ct_y_spacing = ct.PixelSpacing[1]
        pet_x_spacing = pet.PixelSpacing[0]
        pet_y_spacing = pet.PixelSpacing[1]
        ct_intercept = ct.RescaleIntercept
        pet_intercept = pet.RescaleIntercept
        ct_slope = ct.RescaleSlope
        pet_slope = pet.RescaleSlope

        patientAge = int(ct[0x10, 0x1010].value[:-1])
        patientWeight = float(pet[0x10, 0x1030].value)
        patientSize = float(pet[0x10, 0x1020].value)
        Sex = pet[0x10, 0x40].value
        patientSex = 0  # 男性为0， 女性为1
        if (Sex == 'F'):
            patientSex = 1
        DecayFactor = pet.DecayFactor
        a = pet.RadiopharmaceuticalInformationSequence
        TotalDose = a[0].RadionuclideTotalDose
        HalfLife = a[0].RadionuclideHalfLife
        t0 = a[0].RadiopharmaceuticalStartTime
        t1 = pet.AcquisitionTime
        CT_SeriesDescription = line[7]
        PET_SeriesDescription = line[8]
        origin_dir = line[9]
        CT_origin_path = line[10]
        PET_origin_path = line[11]

        infos.append([patientID, z, x, y, r, cancer_type, ct_size, pet_size, ct_x_spacing, ct_y_spacing,
               ct_z_spacing, pet_x_spacing, pet_y_spacing, pet_z_spacing, ct_intercept, pet_intercept,
               ct_slope, pet_slope, patientWeight, patientSex, patientAge,patientSize, TotalDose, t0, t1,
               HalfLife, DecayFactor, CT_SeriesDescription, PET_SeriesDescription, origin_dir,
               CT_origin_path, PET_origin_path])
        print(num, '--->', patientID, 'is ok')
        num = num+1
    df = pd.DataFrame(infos, columns=columns)
    df.to_csv('D:/lung_cancer/data/all_data.csv', index=True)


# 统一all_data.csv文件里的patientSize单位
def unified_patientsize():
    data_path = 'D:/lung_cancer/data/all_data.csv'
    f = csv.reader(open(data_path, 'r'))
    data = []
    for i in f:
        data.append(i)

    new_data = []

    for line in data[1:]:
        Size = float(line[22])
        if Size<10:
            Size = Size*100

        line[22] = int(Size)
        print(line[22])
        new_data.append(line)
    df = pd.DataFrame(new_data, columns=data[0])
    df.to_csv('D:/lung_cancer/data/all_data.csv', index=False)


if __name__ == '__main__':
    # 重新整理柳博的suvmax2.csv
    # rebuild_path()

    # 转存CT和PET目标图像
    # extract_origin_CT_PET()

    # resaved_label()

    # save_info()

    unified_patientsize()